// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/container/f14_map.h>
#include <stdexcept>

#include <pollux/common/base/exceptions.h>
#include <pollux/common/base/simd_util.h>
#include <pollux/vector/simple_vector.h>
#include <pollux/vector/type_aliases.h>

namespace kumo::pollux {

    namespace {
        struct DummyReleaser {
            void addRef() const {}

            void release() const {}
        };
    } // namespace

    template<typename T>
    class ConstantVector final : public SimpleVector<T> {
    public:
        ConstantVector(const ConstantVector &) = delete;

        ConstantVector &operator=(const ConstantVector &) = delete;

        static constexpr bool can_simd =
                (std::is_same_v<T, int64_t> || std::is_same_v<T, int32_t> ||
                 std::is_same_v<T, int16_t> || std::is_same_v<T, int8_t> ||
                 std::is_same_v<T, bool> || std::is_same_v<T, size_t>);

        ConstantVector(
                pollux::memory::MemoryPool *pool,
                size_t length,
                bool isNull,
                TypePtr type,
                T &&val,
                const SimpleVectorStats<T> &stats = {},
                std::optional<ByteCount> represented_bytes = std::nullopt,
                std::optional<ByteCount> storageByteCount = std::nullopt)
                : SimpleVector<T>(
                pool,
                type,
                VectorEncoding::Simple::CONSTANT,
                BufferPtr(nullptr),
                length,
                stats,
                isNull ? 0 : 1 /* distinctValueCount */,
                isNull ? length : 0 /* nullCount */,
                true /*isSorted*/,
                represented_bytes,
                storageByteCount),
                  value_(std::move(val)),
                  isNull_(isNull),
                  initialized_(true) {
            makeNullsBuffer();
            // Special handling for complex types
            if (type->size() > 0) {
                // Only allow null constants to be created through this interface.
                POLLUX_CHECK(isNull_);
                valueVector_ = BaseVector::create(type, 1, pool);
                valueVector_->set_null(0, true);
            }
            if constexpr (std::is_same_v<T, StringView>) {
                if (!isNull_) {
                    // Copy string value.
                    setValue(value_.str());
                }
            }
            // If this is not encoded integer, or string, set value buffer
            if constexpr (can_simd) {
                valueBuffer_ = simd::setAll(value_);
            }
        }

        /// Creates constant vector with value coming from 'index' element of the
        /// 'base' vector. Base vector can be flat or lazy vector. Base vector cannot
        /// be a constant or dictionary vector. Use BaseVector::wrap_in_constant to
        /// automatically peel off encodings of the base vector.
        ///
        /// If base vector is lazy and has not been loaded yet, loading will be
        /// delayed until loaded_vector() is called.
        ConstantVector(
                pollux::memory::MemoryPool *pool,
                vector_size_t length,
                vector_size_t index,
                VectorPtr base,
                const SimpleVectorStats<T> &stats = {})
                : SimpleVector<T>(
                pool,
                base->type(),
                VectorEncoding::Simple::CONSTANT,
                BufferPtr(nullptr),
                length,
                stats,
                std::nullopt,
                std::nullopt,
                true /*isSorted*/,
                base->represented_bytes().has_value()
                ? std::optional<ByteCount>(
                        (base->represented_bytes().value() / base->size()) *
                        length)
                : std::nullopt /* represented_bytes */),
                  valueVector_(base),
                  index_(index) {
            POLLUX_CHECK_NE(
                    base->BaseVector::encoding(),
                    VectorEncoding::Simple::CONSTANT,
                    "Constant vector cannot wrap Constant vector");
            POLLUX_CHECK_NE(
                    base->BaseVector::encoding(),
                    VectorEncoding::Simple::DICTIONARY,
                    "Constant vector cannot wrap Dictionary vector");
            setInternalState();
        }

        virtual ~ConstantVector() override {
            if (auto *wrap_info = wrapInfo_.load()) {
                delete wrap_info;
            }
            if (valueVector_) {
                valueVector_->clear_containing_lazy_and_wrapped();
            }
        }

        bool is_null_at(vector_size_t /*idx*/) const override {
            POLLUX_DCHECK(initialized_);
            return isNull_;
        }

        bool may_have_nulls() const override {
            POLLUX_DCHECK(initialized_);
            return isNull_;
        }

        bool may_have_nulls_recursive() const override {
            POLLUX_DCHECK(initialized_);
            return isNull_ || (valueVector_ && valueVector_->may_have_nulls_recursive());
        }

        void set_null(vector_size_t /*idx*/, bool /*value*/) override {
            POLLUX_FAIL("set_null not supported on ConstantVector");
        }

        const T value() const {
            POLLUX_DCHECK(initialized_);
            return value_;
        }

        const T valueAtFast(vector_size_t /*idx*/) const {
            return value();
        }

        virtual const T value_at(vector_size_t /*idx*/) const override {
            POLLUX_DCHECK(initialized_);
            SimpleVector<T>::checkElementSize();
            return value();
        }

        BufferPtr getStringBuffer() const {
            POLLUX_DCHECK(initialized_);
            return stringBuffer_;
        }

        const T *rawValues() const {
            POLLUX_DCHECK(initialized_);
            return &value_;
        }

        const void *values_as_void() const override {
            POLLUX_DCHECK(initialized_);
            return &value_;
        }

        /// Loads a 256bit vector of data at the virtual byteOffset given
        /// Note this method is implemented on each vector type, but is intentionally
        /// not virtual for performance reasons
        ksimd::batch<T> loadSIMDValueBufferAt(size_t /* byteOffset */) const {
            POLLUX_DCHECK(initialized_);
            return valueBuffer_;
        }

        std::unique_ptr<SimpleVector<uint64_t>> hash_all() const override;

        uint64_t retained_size() const override {
            POLLUX_DCHECK(initialized_);
            if (valueVector_) {
                return valueVector_->retained_size();
            }
            if (stringBuffer_) {
                return stringBuffer_->capacity();
            }
            return sizeof(T);
        }

        BaseVector *loaded_vector() override {
            if (!valueVector_) {
                return this;
            }
            auto loaded = BaseVector::loaded_vector_shared(valueVector_);
            if (loaded == valueVector_) {
                return this;
            }
            valueVector_ = loaded;
            setInternalState();
            return this;
        }

        const BaseVector *loaded_vector() const override {
            return const_cast<ConstantVector<T> *>(this)->loaded_vector();
        }

        bool is_scalar() const override {
            return valueVector_ ? valueVector_->is_scalar()
                                : (this->type_kind() != TypeKind::UNKNOWN);
        }

        const BaseVector *wrapped_vector() const override {
            return valueVector_ ? valueVector_->wrapped_vector() : this;
        }

        vector_size_t wrapped_index(vector_size_t /*index*/) const override {
            return valueVector_ ? valueVector_->wrapped_index(index_) : 0;
        }

        const BufferPtr &wrap_info() const override {
            static const DummyReleaser kDummy;
            auto *wrap_info = wrapInfo_.load();
            if (MELON_UNLIKELY(!wrap_info)) {
                wrap_info = new BufferPtr(BufferView<DummyReleaser>::create(
                        reinterpret_cast<const uint8_t *>(&index_),
                        sizeof(vector_size_t),
                        kDummy));
                BufferPtr *oldWrapInfo = nullptr;
                if (!wrapInfo_.compare_exchange_strong(oldWrapInfo, wrap_info)) {
                    delete wrap_info;
                    wrap_info = oldWrapInfo;
                }
            }
            return *wrap_info;
        }

        /// Base vector if is_scalar() is false (e.g. complex type vector) or if base
        /// vector is a lazy vector that hasn't been loaded yet.
        const VectorPtr &value_vector() const override {
            return valueVector_;
        }

        void set_value_vector(VectorPtr value_vector) override {
            valueVector_ = std::move(value_vector);
        }

        /// Index of the element of the base vector that determines the value of this
        /// constant vector.
        vector_size_t index() const {
            return index_;
        }

        void resize(vector_size_t newSize, bool /*setNotNull*/ = true) override {
            BaseVector::length_ = newSize;
            if constexpr (std::is_same_v<T, StringView>) {
                SimpleVector<StringView>::resizeIsAsciiIfNotEmpty(
                        newSize, SimpleVector<StringView>::getAllIsAscii());
            }
        }

        VectorPtr slice(vector_size_t /*offset*/, vector_size_t length)
        const override {
            POLLUX_DCHECK(initialized_);
            if (valueVector_) {
                return std::make_shared<ConstantVector<T>>(
                        this->pool_, length, index_, valueVector_);
            } else {
                return std::make_shared<ConstantVector<T>>(
                        this->pool_, length, isNull_, this->type_, T(value_));
            }
        }

        void add_nulls(const uint64_t * /*bits*/, const SelectivityVector & /*rows*/)
        override {
            POLLUX_FAIL("add_nulls not supported");
        }

        bool contains_null_at(vector_size_t idx) const override {
            if constexpr (std::is_same_v<T, ComplexType>) {
                if (is_null_at(idx)) {
                    return true;
                }

                return valueVector_->contains_null_at(index_);
            } else {
                return is_null_at(idx);
            }
        }

        void add_nulls(const SelectivityVector & /*rows*/) override {
            POLLUX_FAIL("add_nulls not supported");
        }

        std::optional<int32_t> compare(
                const BaseVector *other,
                vector_size_t index,
                vector_size_t otherIndex,
                CompareFlags flags) const override {
            if constexpr (!std::is_same_v<T, ComplexType>) {
                if (other->is_constant_encoding()) {
                    auto otherConstant = other->as_unchecked<ConstantVector<T>>();
                    if (isNull_ || otherConstant->isNull_) {
                        return BaseVector::compare_nulls(
                                isNull_, otherConstant->isNull_, flags);
                    }

                    auto result = this->typeUsesCustomComparison_
                                  ? SimpleVector<T>::comparePrimitiveAscWithCustomComparison(
                                    this->type_.get(), value_, otherConstant->value_)
                                  : SimpleVector<T>::comparePrimitiveAsc(
                                    value_, otherConstant->value_);

                    return flags.ascending ? result : result * -1;
                }
            }

            return SimpleVector<T>::compare(other, index, otherIndex, flags);
        }

        std::string toString() const {
            if (valueVector_) {
                return valueVector_->toString(index_);
            }

            if (isNull_) {
                return "null";
            } else {
                return SimpleVector<T>::valueToString(value());
            }
        }

        std::string toString(vector_size_t /*index*/) const override {
            return toString();
        }

        bool is_nulls_writable() const override {
            return false;
        }

        void validate(const VectorValidateOptions &options) const override {
            // Do not call BaseVector's validate() since the nulls buffer has
            // a fixed size for constant vectors.
            if (options.callback) {
                options.callback(*this);
            }
            if (valueVector_ != nullptr) {
                if (!isNull_) {
                    POLLUX_CHECK_LT(index_, valueVector_->size());
                }
                valueVector_->validate(options);
            }
        }

        VectorPtr copy_preserve_encodings(
                pollux::memory::MemoryPool *pool = nullptr) const override {
            auto selfPool = pool ? pool : BaseVector::pool_;
            if (valueVector_) {
                return std::make_shared<ConstantVector<T>>(
                        selfPool,
                        BaseVector::length_,
                        index_,
                        valueVector_->copy_preserve_encodings(pool),
                        SimpleVector<T>::stats_);
            }

            return std::make_shared<ConstantVector<T>>(
                    selfPool,
                    BaseVector::length_,
                    isNull_,
                    BaseVector::type_,
                    T(value_),
                    SimpleVector<T>::stats_,
                    BaseVector::representedByteCount_,
                    BaseVector::storageByteCount_);
        }

    protected:
        std::string to_summary_string() const override {
            std::stringstream out;
            out << "[" << BaseVector::encoding() << " "
                << BaseVector::type()->toString() << ": " << BaseVector::size()
                << " elements, " << toString() << "]";

            return out.str();
        }

    private:
        void setInternalState() {
            if (is_lazy_not_loaded(*valueVector_)) {
                POLLUX_CHECK(
                        valueVector_->mark_as_containing_lazy_and_wrapped(),
                        "An unloaded lazy vector cannot be wrapped by two different "
                        "top level vectors.");
                // Do not load Lazy vector
                return;
            }
            // Ensure any internal state in valueVector_ is initialized, and it points
            // to the loaded vector underneath any lazy layers.
            valueVector_ = BaseVector::loaded_vector_shared(valueVector_);

            isNull_ = valueVector_->is_null_at(index_);
            BaseVector::distinctValueCount_ = isNull_ ? 0 : 1;
            const vector_size_t vectorSize = BaseVector::length_;
            BaseVector::nullCount_ = isNull_ ? vectorSize : 0;
            if (valueVector_->is_scalar()) {
                auto simple = valueVector_->loaded_vector()->as<SimpleVector<T>>();
                isNull_ = simple->is_null_at(index_);
                if (!isNull_) {
                    value_ = simple->value_at(index_);
                    if constexpr (std::is_same_v<T, StringView>) {
                        // Copy string value.
                        StringView *valuePtr = reinterpret_cast<StringView *>(&value_);
                        setValue(std::string(valuePtr->data(), valuePtr->size()));

                        auto stringVector = simple->template as<SimpleVector<StringView>>();
                        if (auto ascii = stringVector->isAscii(index_)) {
                            SimpleVector<T>::setAllIsAscii(ascii.value());
                        }
                    }
                }
                valueVector_ = nullptr;
            }
            makeNullsBuffer();
            initialized_ = true;
        }

        void setValue(const std::string &string) {
            BaseVector::in_memory_bytes_ += string.size();
            value_ = pollux::to<decltype(value_)>(string);
            if constexpr (can_simd) {
                valueBuffer_ = simd::setAll(value_);
            }
        }

        void makeNullsBuffer() {
            if (!isNull_) {
                return;
            }
            BaseVector::nulls_ =
                    AlignedBuffer::allocate<uint64_t>(1, BaseVector::pool());
            BaseVector::nulls_->setSize(1);
            BaseVector::rawNulls_ = BaseVector::nulls_->as<uint64_t>();
            *BaseVector::nulls_->asMutable<uint64_t>() = bits::kNull64;
        }

        // 'valueVector_' element 'index_' represents a complex constant
        // value. 'valueVector_' is nullptr if the constant is scalar.
        VectorPtr valueVector_;
        // The index of the represented value in 'valueVector_'.
        vector_size_t index_ = 0;
        // Holds the memory for backing non-inlined values represented by StringView.
        BufferPtr stringBuffer_;
        T value_;
        bool isNull_ = false;
        bool initialized_{false};
        mutable std::atomic<BufferPtr *> wrapInfo_{nullptr};

        // This must be at end to avoid memory corruption.
        std::conditional_t<can_simd, ksimd::batch<T>, char> valueBuffer_;
    };

    template<>
    void ConstantVector<StringView>::setValue(const std::string &string);

    template<>
    void ConstantVector<std::shared_ptr<void>>::setValue(const std::string &string);

    template<>
    void ConstantVector<ComplexType>::setValue(const std::string &string);

    template<typename T>
    using ConstantVectorPtr = std::shared_ptr<ConstantVector<T>>;

    template<typename T>
    std::unique_ptr<SimpleVector<uint64_t>> ConstantVector<T>::hash_all() const {
        return std::make_unique<ConstantVector<uint64_t>>(
                BaseVector::pool_,
                BaseVector::length_,
                false /* isNull */,
                BIGINT(),
                this->hash_value_at(0),
                SimpleVectorStats<uint64_t>{}, /* stats */
                sizeof(uint64_t) * BaseVector::length_ /* represented_bytes */);
    }

} // namespace kumo::pollux

